#!/bin/bash

# This script is used to install proprietary drivers for LicheePi 4A (TH1520)
# Author: misaka00251 <liuxin@iscas.ac.cn>
# Copyright (C) 2024 misaka00251

# Some colors for the output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Functions here

# ========== Install proprietary drivers ==========

# We need to remove some packages
pre_remove() {
    echo -e "Some cleanup is needed before installing the proprietary drivers..."
    dnf remove xorg-x11-drv-qxl -y
    # Check if the command was successful
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}\nCleanup success. Proceeding to the next step...${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        exit 1
    fi
}

# Then, we need to change repositories
update_hardware_repo_url() {
    local repo_file="/etc/yum.repos.d/hw.repo"
    local new_url="https://build-repo.tarsier-infra.isrc.ac.cn/Factory:/Board:/TH1520:/Script/24.03/"
    local priority_value="50"

    if [ -f "$repo_file" ]; then
        # Update the baseurl
        sed -i "s|^baseurl=.*|baseurl=${new_url}|" "$repo_file"

        # Ensure priority is 50
        if grep -q "^priority=" "$repo_file"; then
            sed -i "s|^priority=.*|priority=${priority_value}|" "$repo_file"
        else
            echo "priority=${priority_value}" >> "$repo_file"
        fi

        echo -e "${GREEN}\nRepo convert success. Proceeding to the next step...${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        return 1
    fi
}

# Then do a dnf update
do_update() {
    echo -e "\nUpdating the system...\n"
    dnf clean all
    sleep 3
    dnf update -y
    # Check if the command was successful
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}\nSystem update success. Proceeding to the next step...${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        exit 1
    fi
}

# Then we need to install extra packages
install_extra_packages() {
    echo -e "\nInstalling extra packages & proprietary drivers...\n"
    dnf install thead-gles-addons-6 libdrm glmark2 xfwm4 NetworkManager-wifi th1520-bluetooth th1520-vpu th1520-npu mesa-vdpau-drivers gstreamer1-plugins-base-1.20.4 -y
    # Check if the command was successful
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}\nExtra packages & proprietary drivers installed successfully!${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        exit 1
    fi
}

# ========== Uninstall proprietary drivers ==========

# First, we need to remove the extra packages & proprietary drivers
# These packages are not available in the official openEuler repository
# So, we need to remove them manually
uninstall_extra_packages() {
    echo -e "\nUninstalling extra packages & proprietary drivers...\n"
    # Can't remove libdrm here
    # TODO: NetworkManager-wifi & mesa-vdpau-drivers?
    dnf remove thead-gles-addons-6 glmark2 th1520-bluetooth th1520-vpu th1520-npu thead-gles-addons-6 NetworkManager-wifi mesa-vdpau-drivers --noautoremove -y
    # gstreamer1-plugins-base-1.20.4 can be uninstalled by `dnf update --disablerepo hw`
    # Because it's an downgrade package
    dnf update --disablerepo hw -y
    # Check if the command was successful
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}\nExtra packages & proprietary drivers uninstalled successfully!${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        exit 1
    fi
}

# Then, change the repository priority to 100
# So that the official openEuler repository is used everytime (which is 90)
update_official_repo_url_back() {
    local repo_file="/etc/yum.repos.d/hw.repo"
    local priority_value="100"

    if [ -f "$repo_file" ]; then
        # Ensure priority is 100
        if grep -q "^priority=" "$repo_file"; then
            sed -i "s|^priority=.*|priority=${priority_value}|" "$repo_file"
        else
            echo "priority=${priority_value}" >> "$repo_file"
        fi

        echo -e "${GREEN}\nRepo convert success. Proceeding to the next step...${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        return 1
    fi
    # Then do a dnf update
    dnf clean all
    sleep 3
    dnf update -y
}


# Finally, we need to downgrade any packages that were updated
# Like mesa & xorg family. We need to downgrade them to the official openEuler repository
# This is to ensure the system is stable
do_downgrade() {
    echo -e "\nDowngrading packages...\n"
    # First, downgrade mesa
    dnf downgrade mesa* -y
    # Then, downgrade xorg-x11*
    dnf downgrade xorg-x11-server -y
    dnf reinstall xorg* --disablerepo hw -y
    # Check if the command was successful
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}\nPackages downgraded successfully!${NC}\n"
    else
        echo -e "${RED}\nSomething is wrong. Please try again or contact your system admin.${NC}\n"
        exit 1
    fi
}

# Show our oerv ascii art at every launch

echo -e "${BLUE}                     .:*##*-.                     ${NC}"
echo -e "${BLUE}                 ..-#%%%##%%%#-..                 ${NC}"
echo -e "${BLUE}               .+#%%%#-....-#%%%#+.               ${NC}"
echo -e "${BLUE}          ..:+%%%%*:.         :*%%%%+:..          ${NC}"
echo -e "${BLUE}        .-*%%%%=:.              .:=%%%%#-.        ${NC}"
echo -e "${BLUE}    ..=#%%%#=.                      .=*%%%#=:.    ${NC}"
echo -e "${BLUE}   :#%%%*:..  ${CYAN}.      ...              ${BLUE}..:*%%%#:   ${NC}"
echo -e "${BLUE}   +%%:.     ${CYAN}.*%%%%%%%%%#=.               ${BLUE}.:%%+   ${NC}"
echo -e "${BLUE}   *%%.      ${CYAN}:#%%%***#%%%%*.               ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.     ${CYAN}.+%%%-   .=%%%%:               ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.     ${CYAN}:#%%*--==+#%%%*.               ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.     ${CYAN}-%%%###%%%%%#+.                ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.    ${CYAN}.#%%*. .:*%%%*:.          ${YELLOW}.     ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.   ${CYAN}.=%%%:     =#%%#.${YELLOW}:-=========-:.  ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.   ${CYAN}.#%%*.     .-#%%#:${YELLOW}......-==-...  ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.  ${CYAN}.=###:        .::::.   ${YELLOW}.-==-.     ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.                 ${YELLOW}.-==-. .-==-.      ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   *%%.                  ${YELLOW}.-===:-=-:        ${BLUE}.%%*   ${NC}"
echo -e "${BLUE}   +%%:.                  ${YELLOW}.:-====:.       ${BLUE}.:%%+   ${NC}"
echo -e "${BLUE}   :#%%%*:..                ${YELLOW}.-=-.     ${BLUE}..:*%%%#:   ${NC}"
echo -e "${BLUE}    ..=#%%%#=.               ${YELLOW}.:.   ${BLUE}..=*%%%#=:.    ${NC}"
echo -e "${BLUE}        .-*%%%%=:.              .:=%%%%#-.        ${NC}"
echo -e "${BLUE}          ..:+%%%%*:.         :*%%%%+:..          ${NC}"
echo -e "${BLUE}               .+#%%%#-....-#%%%#+.               ${NC}"
echo -e "${BLUE}                 ..-#%%%##%%%#-..                 ${NC}"
echo -e "${BLUE}                     .:*##*-.                     ${NC}"
echo -e "${CYAN}                Version:  20240902                ${NC}"

echo -e "${NC}\n\n                 openEuler RISC-V                 "
echo -e "                     Presents                     \n\n"


# Load OS information
. /etc/os-release

# Check if the OS is openEuler
if [ "$NAME" != "openEuler" ]; then
    echo -e "${RED}This script is only supported on openEuler.${NC}"
    exit 1
fi

# Check if the OS version is 24.03
if [ "$VERSION_ID" != "24.03" ]; then
    echo -e "${RED}This script is only supported on openEuler 24.03 (LTS).${NC}"
    exit 1
fi

# Check if you are on riscv64 architecture
if [ "$(uname -m)" != "riscv64" ]; then
    echo -e "${RED}This script is only supported on riscv64 architecture.${NC}"
    exit 1
fi

# You can't run this script directly
# You need to append parameters to the script
if [ $# -eq 0 ]; then
    echo -e "${RED}You can't run this script directly!${NC}"
    echo -e "You need to append parameters to the script!"
    echo -e "For more information, please run:"
    echo -e "${GREEN}bash th1520-proprietary.sh help${NC}"
    exit 1
fi

# Help menu
if [ $1 == "help" ]; then
    echo -e "${GREEN}Usage:${NC}"
    echo -e "bash th1520-proprietary.sh [OPTION]"
    echo -e ""
    echo -e "${GREEN}Options:${NC}"
    echo -e "install - Install proprietary drivers for LicheePi 4A (TH1520)"
    echo -e "uninstall - Uninstall proprietary drivers for LicheePi 4A (TH1520)"
    exit 0
fi

# Install proprietary drivers
if [ $1 == "install" ]; then
    echo -e "${GREEN}Installing proprietary drivers for LicheePi 4A (TH1520)...${NC}"
    # Check if the user is root
    if [ "$EUID" -ne 0 ]; then
        echo -e "${RED}Please run this script as root!${NC}"
        exit 1
    fi
    echo -e "${YELLOW}Please wait...${NC}"
    pre_remove
    update_hardware_repo_url
    do_update
    install_extra_packages
    echo -e "${GREEN}Proprietary drivers for LicheePi 4A (TH1520) have been installed successfully!${NC}"
    echo -e "${GREEN}Please reboot your system to apply the changes.${NC}"
    exit 0
fi

# Uninstall proprietary drivers
if [ $1 == "uninstall" ]; then
    echo -e "${GREEN}Uninstalling proprietary drivers for LicheePi 4A (TH1520)...${NC}"
    # Check if the user is root
    if [ "$EUID" -ne 0 ]; then
        echo -e "${RED}Please run this script as root!${NC}"
        exit 1
    fi
    echo -e "${YELLOW}Please wait...${NC}"
    uninstall_extra_packages
    update_official_repo_url_back
    do_downgrade
    # Check if the command was successful
    echo -e "${GREEN}Proprietary drivers for LicheePi 4A (TH1520) have been uninstalled successfully!${NC}"
    echo -e "${GREEN}Please reboot your system to apply the changes.${NC}"
    exit 0
fi
